{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Attractivity vs Stability"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the underactuated repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  Colab will ask you to \"Reset all runtimes\"; say no to save yourself the reinstall.\n",
    "- On Binder, the machines should already be provisioned by the time you can run this; it should return (almost) instantly.\n",
    "\n",
    "More details are available [here](http://underactuated.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import pydrake\n",
    "    import underactuated\n",
    "except ImportError:\n",
    "    !curl -s https://raw.githubusercontent.com/RussTedrake/underactuated/master/scripts/setup/jupyter_setup.py > jupyter_setup.py\n",
    "    from jupyter_setup import setup_underactuated\n",
    "    setup_underactuated()\n",
    "\n",
    "# Setup matplotlib.  \n",
    "from IPython import get_ipython\n",
    "if get_ipython() is not None: get_ipython().run_line_magic(\"matplotlib\", \"inline\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# python libraries\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# pydrake imports\n",
    "from pydrake.all import DiagramBuilder, Variable, SymbolicVectorSystem, LogOutput, Simulator\n",
    "\n",
    "# underactuated imports\n",
    "from underactuated import plot_2d_phase_portrait"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dynamics of the Nonlinear System\n",
    "Here we create a Drake dynamical system using the class `SymbolicVectorSystem`.\n",
    "This requires the dynamics to be written in the form $\\dot{\\mathbf{x}} = f(\\mathbf{x})$, where $\\mathbf{x}$ is a vector of Drake symbolic variables.\n",
    "`SymbolicVectorSystem` is just one of the many options we have in Drake: when systems will get more complicated, writing the dynamics by hand can be rather tedious..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function that given the current state returns the state time derivative\n",
    "def dynamics(x):\n",
    "    r = np.sqrt(x[0]**2 + x[1]**2)\n",
    "    return [\n",
    "        x[0]*(1-r) - x[1]*(r-x[0])/(2*r),\n",
    "        x[1]*(1-r) + x[0]*(r-x[0])/(2*r)\n",
    "    ]\n",
    "\n",
    "# state variables\n",
    "x1 = Variable(\"x1\")\n",
    "x2 = Variable(\"x2\")\n",
    "x = [x1, x2]\n",
    "\n",
    "# Drake nonlinear system\n",
    "system = SymbolicVectorSystem(state=x, output=x,  dynamics=dynamics(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Drake Diagram\n",
    "We then construct a Drake diagram.\n",
    "This is nothing more than a set of interconnected dynamical systems (similiar to the Simulink idea, if you ever used it).\n",
    "Our diagram will be very simple: we just connect our dynamical system to a logger, which will measure and store the system state during the simulation (similar to the Simulink `To Workspace` block)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize builder of the diagram\n",
    "builder = DiagramBuilder()\n",
    "\n",
    "# add our dynamical system\n",
    "# (note: builder.AddSystem() returns a pointer to the system passed as input,\n",
    "# hence it is safe to assign the name \"system\" to its output)\n",
    "system = builder.AddSystem(system)\n",
    "\n",
    "# logger block to measure and store the state\n",
    "# connected to the (first and only) output port of the dynamical system\n",
    "logger = LogOutput(system.get_output_port(0), builder)\n",
    "\n",
    "# finalize diagram\n",
    "diagram = builder.Build()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulation\n",
    "We are ready to simulate our dynamical system.\n",
    "To this end we just feed our diagram in a the Drake `Simulator` and `AdvanceTo` the desired time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function that given the initial state\n",
    "# and a simulation time returns the system trajectory\n",
    "def simulate(x, sim_time):\n",
    "    \n",
    "    # clean the logger from old trajectories\n",
    "    logger.reset()\n",
    "    \n",
    "    # set up the simulator\n",
    "    simulator = Simulator(diagram)\n",
    "    \n",
    "    # set initial conditions\n",
    "    # (for now, think of \"context\" as a synonym of state)\n",
    "    context = simulator.get_mutable_context()\n",
    "    context.SetContinuousState(x)\n",
    "    \n",
    "    # simulate from t=0 to t=sim_time\n",
    "    simulator.AdvanceTo(sim_time)\n",
    "    \n",
    "    # return the output (here = state) trajectory\n",
    "    return logger.data()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plotting Helper\n",
    "Let us write a simple function that plots the system trajectory given the initial conditions and the duration of the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function that given the initial state\n",
    "# and a simulation time plots the system trajectory\n",
    "def plot_trajectory(x, sim_time):\n",
    "    \n",
    "    # simulate the system from the given state\n",
    "    traj = simulate(x, sim_time)\n",
    "    \n",
    "    # plot a blue dot for the initial conditions\n",
    "    label = r'$\\mathbf{x}(0)=[%.2f,%.2f]^T$'%(x[0], x[1])\n",
    "    plt.scatter(*x, s=50, c='b', zorder=3, label=label)\n",
    "    \n",
    "    # plot a red curve for the trajectory\n",
    "    label = r'$\\mathbf{x}(t)$'\n",
    "    plt.plot(traj[0,:], traj[1,:], label=label, c='r')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Play with the Initial Conditions\n",
    "Now we can finally visualize the dynamics of our system.\n",
    "To do this, in the next cell modify the two variables (currently set to arbitrary values):\n",
    "- `initial_conditions`: state of the system a time zero,\n",
    "- `sim_time`: duration of the simulation in seconds.\n",
    "\n",
    "Then, run the last cell to see the result of the simulation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_conditions = [-1, -1] # modify here\n",
    "sim_time = 1 # modify here"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize plot and make it big enough\n",
    "plt.figure(figsize=(10, 10))\n",
    "\n",
    "# plot the phase portrait of the 2d system\n",
    "plot_2d_phase_portrait(dynamics, x1lim=[-2, 2], x2lim=[-2, 2], linewidth=1, density=2)\n",
    "\n",
    "# superimpose the trajectory to the phase portrait\n",
    "plot_trajectory(initial_conditions, sim_time)\n",
    "\n",
    "# add legend\n",
    "plt.legend(loc=1)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}